#ifndef TEST_CNN_BGV_MNIST_HELPER
#define TEST_CNN_BGV_MNIST_HELPER

#include "CNN/mnist_helper.hpp"
#include "hypercube/hypercube_cryto_agent.hpp"
#include "hypercube/hypercube_hmm_engine.hpp"

#define T_NUM 4

static constexpr double SCALE = 1024.0;

static constexpr unsigned long MATRIX_DIM_M = 64;
static constexpr unsigned long MATRIX_DIM_N = 64;
static constexpr unsigned long MATRIX_DIM_K = 64;
static constexpr unsigned long HYPERCUBE_DIM_X =
    MAX(MATRIX_DIM_M, MATRIX_DIM_K);
static constexpr unsigned long HYPERCUBE_DIM_Y =
    MAX(MATRIX_DIM_K, MATRIX_DIM_N);

static cube_signature MNIST_HYPERCUBE({HYPERCUBE_DIM_X, HYPERCUBE_DIM_Y});

class bgv_hmnist_client : public h_mnist_client<helib::BGV> {
private:
  bgv_hmnist_client() = delete;

  helib::Ctxt init_hx(const torch::Tensor &a, int asz0, int t_index) const;

  helib::Ctxt init_hconv(const torch::Tensor &a, int asz0, int asz1,
                         int t_index) const;

  helib::Ctxt init_hfc1(const torch::Tensor &a, int asz0, int t_index) const;

  helib::Ctxt init_hfc2(const torch::Tensor &a, int t_index) const;

public:
  bgv_hmnist_client(const params<helib::BGV> &params, int crt_num);
  ~bgv_hmnist_client() {
    for (auto hmmcc : hmmccs_) {
      delete hmmcc;
    }
  }

  int get_crt_num() const { return crt_num_; }

  const std::vector<helib::Context *> &get_context() const { return contexts_; }

  // Generate encrypted model
  void init_hcnn(encrypted_model<helib::BGV> &ctxt_model,
                 const torch::nn::Conv2d &conv1, const torch::nn::Linear &fc1,
                 const torch::nn::Linear &fc2) const override;

  // Generate encrypted image
  void init_input(encrypted_input<helib::BGV> &ctxt_input,
                  const torch::Tensor &filter,
                  const torch::Tensor &x) const override;

  void recover_result(at::Tensor &out,
                      const encrypted_output<helib::BGV> &in) const override;

private:
  ENC key = ENC::PUBKEY;

  int crt_num_;
  NTL::ZZ T_;
  std::vector<NTL::ZZ> ti_;

  std::vector<hypercube_hmmcc *> hmmccs_;
  std::vector<helib::Context *> contexts_;
};

class bgv_hmnist_server : public h_mnist_server<helib::BGV> {
private:
  bgv_hmnist_server() = delete;

public:
  bgv_hmnist_server(const std::vector<helib::Context *> &contexts,
                    const hmm_status<shmm_engine> &status, int crt_num);
  ~bgv_hmnist_server() {
    for (auto hmme : hmmes_) {
      delete hmme;
    }
  }

  void
  forward_test(encrypted_output<helib::BGV> &out,
               const encrypted_input<helib::BGV> &in,
               const encrypted_model<helib::BGV> &ctxt_model) const override;

  void report_metrics(std::ostream &ss = std::cout) const override;

private:
  std::vector<metric> metrics_;

  int crt_num_;

  std::vector<hypercube_hmme *> hmmes_;
};

#endif /* TEST_CNN_BGV_MNIST_HELPER */
